"""AES-256-GCM encryption for API keys."""

import os
import base64
import hashlib
import hmac
from typing import Tuple, Optional
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.backends import default_backend


class KeyEncryptor:
    """AES-256-GCM encryption for secure API key storage."""
    
    # Security constants
    KEY_SIZE = 32  # 256 bits
    SALT_SIZE = 32  # 256 bits
    NONCE_SIZE = 12  # 96 bits for GCM
    TAG_SIZE = 16  # 128 bits
    ITERATIONS = 100000  # PBKDF2 iterations
    
    def __init__(self, master_key: Optional[str] = None):
        """Initialize encryptor with master key.
        
        Args:
            master_key: Master encryption key. If not provided, reads from MASTER_KEY env var.
        """
        if master_key:
            self._master_key = master_key.encode()
        else:
            env_key = os.environ.get("MASTER_KEY")
            if not env_key:
                raise ValueError("Master key not provided and MASTER_KEY environment variable not set")
            self._master_key = env_key.encode()
            
        # Validate master key strength
        if len(self._master_key) < 32:
            raise ValueError("Master key must be at least 32 characters long")
            
    def derive_key(self, salt: bytes, context: bytes = b"api_key_encryption") -> bytes:
        """Derive encryption key from master key using PBKDF2.
        
        Args:
            salt: Random salt for key derivation
            context: Additional context for key derivation
            
        Returns:
            Derived 256-bit key
        """
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=self.KEY_SIZE,
            salt=salt + context,
            iterations=self.ITERATIONS,
            backend=default_backend()
        )
        return kdf.derive(self._master_key)
    
    def encrypt(self, plaintext: str) -> Tuple[str, str]:
        """Encrypt API key using AES-256-GCM.
        
        Args:
            plaintext: Plain text API key
            
        Returns:
            Tuple of (encrypted_key_base64, key_hash)
        """
        # Generate random salt and nonce
        salt = os.urandom(self.SALT_SIZE)
        nonce = os.urandom(self.NONCE_SIZE)
        
        # Derive encryption key
        key = self.derive_key(salt)
        
        # Create cipher
        cipher = Cipher(
            algorithms.AES(key),
            modes.GCM(nonce),
            backend=default_backend()
        )
        encryptor = cipher.encryptor()
        
        # Encrypt the plaintext
        plaintext_bytes = plaintext.encode('utf-8')
        ciphertext = encryptor.update(plaintext_bytes) + encryptor.finalize()
        
        # Get authentication tag
        tag = encryptor.tag
        
        # Combine salt + nonce + tag + ciphertext
        encrypted_data = salt + nonce + tag + ciphertext
        
        # Base64 encode for storage
        encrypted_b64 = base64.b64encode(encrypted_data).decode('utf-8')
        
        # Generate key hash for duplicate detection
        key_hash = self.generate_key_hash(plaintext)
        
        return encrypted_b64, key_hash
    
    def decrypt(self, encrypted_b64: str) -> str:
        """Decrypt API key from storage.
        
        Args:
            encrypted_b64: Base64 encoded encrypted data
            
        Returns:
            Decrypted API key
            
        Raises:
            ValueError: If decryption fails (invalid key or corrupted data)
        """
        try:
            # Decode from base64
            encrypted_data = base64.b64decode(encrypted_b64)
            
            # Extract components
            salt = encrypted_data[:self.SALT_SIZE]
            nonce = encrypted_data[self.SALT_SIZE:self.SALT_SIZE + self.NONCE_SIZE]
            tag = encrypted_data[self.SALT_SIZE + self.NONCE_SIZE:self.SALT_SIZE + self.NONCE_SIZE + self.TAG_SIZE]
            ciphertext = encrypted_data[self.SALT_SIZE + self.NONCE_SIZE + self.TAG_SIZE:]
            
            # Derive decryption key
            key = self.derive_key(salt)
            
            # Create cipher
            cipher = Cipher(
                algorithms.AES(key),
                modes.GCM(nonce, tag),
                backend=default_backend()
            )
            decryptor = cipher.decryptor()
            
            # Decrypt
            plaintext = decryptor.update(ciphertext) + decryptor.finalize()
            
            return plaintext.decode('utf-8')
            
        except Exception as e:
            raise ValueError(f"Decryption failed: {str(e)}")
    
    def generate_key_hash(self, api_key: str) -> str:
        """Generate secure hash of API key for duplicate detection.
        
        Args:
            api_key: Plain text API key
            
        Returns:
            Hex-encoded hash of the key
        """
        # Use HMAC with master key to prevent rainbow table attacks
        h = hmac.new(
            self._master_key,
            api_key.encode('utf-8'),
            hashlib.sha256
        )
        return h.hexdigest()
    
    def verify_integrity(self, encrypted_b64: str, expected_hash: str) -> bool:
        """Verify integrity of encrypted key.
        
        Args:
            encrypted_b64: Encrypted key data
            expected_hash: Expected hash of the original key
            
        Returns:
            True if integrity is verified, False otherwise
        """
        try:
            # Decrypt the key
            decrypted = self.decrypt(encrypted_b64)
            
            # Generate hash of decrypted key
            actual_hash = self.generate_key_hash(decrypted)
            
            # Constant-time comparison to prevent timing attacks
            return hmac.compare_digest(actual_hash, expected_hash)
            
        except Exception:
            return False
    
    def rotate_encryption(self, old_encrypted_b64: str, new_master_key: str) -> Tuple[str, str]:
        """Rotate encryption with new master key.
        
        Args:
            old_encrypted_b64: Data encrypted with current master key
            new_master_key: New master key for encryption
            
        Returns:
            Tuple of (new_encrypted_key_base64, new_key_hash)
        """
        # Decrypt with current key
        plaintext = self.decrypt(old_encrypted_b64)
        
        # Create new encryptor with new master key
        new_encryptor = KeyEncryptor(new_master_key)
        
        # Encrypt with new key
        return new_encryptor.encrypt(plaintext)


class SecureString:
    """Secure string wrapper that clears memory after use."""
    
    def __init__(self, value: str):
        """Initialize secure string.
        
        Args:
            value: String value to secure
        """
        self._value = value
        
    def get(self) -> str:
        """Get the string value."""
        return self._value
    
    def clear(self):
        """Clear the string from memory."""
        if self._value:
            # Overwrite memory with zeros
            value_len = len(self._value)
            self._value = '\x00' * value_len
            self._value = None
            
    def __del__(self):
        """Ensure cleanup on deletion."""
        self.clear()
        
    def __enter__(self):
        """Context manager entry."""
        return self
        
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - clear memory."""
        self.clear()


def validate_key_strength(api_key: str) -> Tuple[bool, Optional[str]]:
    """Validate API key strength and format.
    
    Args:
        api_key: API key to validate
        
    Returns:
        Tuple of (is_valid, error_message)
    """
    # Check minimum length
    if len(api_key) < 20:
        return False, "API key is too short (minimum 20 characters)"
        
    # Check for common test keys
    test_patterns = ['test', 'demo', 'example', '1234', '0000']
    key_lower = api_key.lower()
    for pattern in test_patterns:
        if pattern in key_lower:
            return False, f"API key appears to be a test key (contains '{pattern}')"
            
    # Check for sufficient entropy (basic check)
    unique_chars = len(set(api_key))
    if unique_chars < 10:
        return False, "API key has insufficient entropy"
        
    return True, None